#!/usr/bin/env python3
"""
Intel HEX File Parser
支持解析完整的HEX文件和单行HEX记录
"""

import sys
import os
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from enum import Enum


class HexRecordType(Enum):
    """HEX记录类型枚举"""
    DATA = 0x00                    # 数据记录
    END_OF_FILE = 0x01            # 文件结束记录
    EXTENDED_SEGMENT_ADDRESS = 0x02 # 扩展段地址记录
    START_SEGMENT_ADDRESS = 0x03   # 开始段地址记录
    EXTENDED_LINEAR_ADDRESS = 0x04 # 扩展线性地址记录
    START_LINEAR_ADDRESS = 0x05    # 开始线性地址记录


@dataclass
class HexRecord:
    """HEX记录结构"""
    length: int
    address: int
    record_type: HexRecordType
    data: List[int]
    checksum: int
    valid: bool = False


@dataclass
class SegmentInfo:
    """段信息结构"""
    start_address: int
    end_address: int
    data_size: int
    data: List[int]


@dataclass
class HexFileInfo:
    """HEX文件信息结构"""
    filename: str
    total_records: int = 0
    data_records: int = 0
    total_data_bytes: int = 0
    start_address: int = 0
    end_address: int = 0
    extended_linear_address: int = 0
    extended_segment_address: int = 0
    memory_map: Dict[int, int] = None
    records: List[HexRecord] = None
    segments: List[SegmentInfo] = None
    
    def __post_init__(self):
        if self.memory_map is None:
            self.memory_map = {}
        if self.records is None:
            self.records = []
        if self.segments is None:
            self.segments = []


class HexParser:
    """HEX文件解析器"""
    
    def __init__(self):
        self.file_info = HexFileInfo("")
    
    def parse_file(self, filename: str) -> bool:
        """解析HEX文件"""
        try:
            with open(filename, 'r') as file:
                lines = file.readlines()
        except IOError as e:
            print(f"Error: Cannot open file {filename}")
            return False
        
        self.file_info = HexFileInfo(filename)
        self.file_info.records.clear()
        self.file_info.memory_map.clear()
        
        current_address = 0
        base_address = 0
        
        for line in lines:
            line = line.strip()
            if not line:
                continue
            
            if not line.startswith(':'):
                print(f"Warning: Skipping invalid line: {line}")
                continue
            
            record = self._parse_record(line)
            self.file_info.records.append(record)
            self.file_info.total_records += 1
            
            if not record.valid:
                print(f"Error: Invalid record: {line}")
                continue
            
            # 处理不同类型的记录
            if record.record_type == HexRecordType.DATA:
                self.file_info.data_records += 1
                self.file_info.total_data_bytes += record.length
                
                # 计算实际地址
                current_address = base_address + record.address
                
                # 将数据添加到内存映射
                for i, byte_data in enumerate(record.data):
                    self.file_info.memory_map[current_address + i] = byte_data
                
                # 更新地址范围
                if self.file_info.start_address == 0 or current_address < self.file_info.start_address:
                    self.file_info.start_address = current_address
                if current_address + record.length > self.file_info.end_address:
                    self.file_info.end_address = current_address + record.length - 1
                    
            elif record.record_type == HexRecordType.EXTENDED_LINEAR_ADDRESS:
                if len(record.data) >= 2:
                    self.file_info.extended_linear_address = (record.data[0] << 8) | record.data[1]
                    base_address = self.file_info.extended_linear_address << 16
                    
            elif record.record_type == HexRecordType.EXTENDED_SEGMENT_ADDRESS:
                if len(record.data) >= 2:
                    self.file_info.extended_segment_address = (record.data[0] << 8) | record.data[1]
                    base_address = self.file_info.extended_segment_address << 4
                    
            elif record.record_type == HexRecordType.START_LINEAR_ADDRESS:
                if len(record.data) >= 4:
                    self.file_info.start_address = (
                        (record.data[0] << 24) |
                        (record.data[1] << 16) |
                        (record.data[2] << 8) |
                        record.data[3]
                    )
                    
            elif record.record_type == HexRecordType.END_OF_FILE:
                print("End of file record detected")
        
        return True
    
    def parse_single_record(self, record: str) -> bool:
        """解析单行HEX记录"""
        # 清空之前的数据
        self.file_info = HexFileInfo("Single HEX Record")
        
        # 解析单行记录
        parsed_record = self._parse_record(record)
        self.file_info.records.append(parsed_record)
        self.file_info.total_records = 1
        
        if not parsed_record.valid:
            print("Error: Invalid HEX record format")
            return False
        
        # 处理记录类型
        if parsed_record.record_type == HexRecordType.DATA:
            self.file_info.data_records = 1
            self.file_info.total_data_bytes = parsed_record.length
            
            # 将数据添加到内存映射
            for i, byte_data in enumerate(parsed_record.data):
                self.file_info.memory_map[parsed_record.address + i] = byte_data
            
            self.file_info.start_address = parsed_record.address
            self.file_info.end_address = parsed_record.address + parsed_record.length - 1
            
        elif parsed_record.record_type == HexRecordType.START_LINEAR_ADDRESS:
            if len(parsed_record.data) >= 4:
                self.file_info.start_address = (
                    (parsed_record.data[0] << 24) |
                    (parsed_record.data[1] << 16) |
                    (parsed_record.data[2] << 8) |
                    parsed_record.data[3]
                )
                
        elif parsed_record.record_type == HexRecordType.EXTENDED_LINEAR_ADDRESS:
            if len(parsed_record.data) >= 2:
                self.file_info.extended_linear_address = (parsed_record.data[0] << 8) | parsed_record.data[1]
                
        elif parsed_record.record_type == HexRecordType.EXTENDED_SEGMENT_ADDRESS:
            if len(parsed_record.data) >= 2:
                self.file_info.extended_segment_address = (parsed_record.data[0] << 8) | parsed_record.data[1]
        
        return True
    
    def _parse_record(self, line: str) -> HexRecord:
        """解析单行HEX记录"""
        if len(line) < 11 or not line.startswith(':'):
            return HexRecord(0, 0, HexRecordType.DATA, [], 0, False)
        
        try:
            # 解析长度
            length = int(line[1:3], 16)
            
            # 解析地址
            address = int(line[3:7], 16)
            
            # 解析记录类型
            record_type = HexRecordType(int(line[7:9], 16))
            
            # 解析数据
            data_start = 9
            data_length = length * 2
            if len(line) < data_start + data_length + 2:
                return HexRecord(0, 0, HexRecordType.DATA, [], 0, False)
            
            data_str = line[data_start:data_start + data_length]
            data = [int(data_str[i:i+2], 16) for i in range(0, len(data_str), 2)]
            
            # 解析校验和
            checksum = int(line[data_start + data_length:data_start + data_length + 2], 16)
            
            # 创建记录
            record = HexRecord(length, address, record_type, data, checksum)
            
            # 验证校验和
            record.valid = self._validate_checksum(record)
            
            return record
            
        except (ValueError, IndexError) as e:
            print(f"Error parsing record: {e}")
            return HexRecord(0, 0, HexRecordType.DATA, [], 0, False)
    
    def _validate_checksum(self, record: HexRecord) -> bool:
        """验证校验和"""
        calculated_checksum = self._calculate_checksum(record)
        return calculated_checksum == record.checksum
    
    def _calculate_checksum(self, record: HexRecord) -> int:
        """计算校验和"""
        checksum = record.length
        checksum += (record.address >> 8) & 0xFF
        checksum += record.address & 0xFF
        checksum += record.record_type.value
        
        for byte_data in record.data:
            checksum += byte_data
        
        return (~checksum + 1) & 0xFF
    
    def analyze_segments(self):
        """分析段信息"""
        self.file_info.segments.clear()
        
        if not self.file_info.memory_map:
            return
        
        # 按地址排序找到连续的段
        sorted_data = sorted(self.file_info.memory_map.items())
        
        if not sorted_data:
            return
        
        current_segment = SegmentInfo(sorted_data[0][0], 0, 0, [sorted_data[0][1]])
        
        for i in range(1, len(sorted_data)):
            current_addr = sorted_data[i][0]
            prev_addr = sorted_data[i-1][0]
            
            # 检查是否连续（允许小的间隔）
            if current_addr - prev_addr <= 1:
                # 连续或相邻，添加到当前段
                current_segment.data.append(sorted_data[i][1])
            else:
                # 不连续，结束当前段，开始新段
                current_segment.end_address = prev_addr
                current_segment.data_size = len(current_segment.data)
                self.file_info.segments.append(current_segment)
                
                # 开始新段
                current_segment = SegmentInfo(current_addr, 0, 0, [sorted_data[i][1]])
        
        # 添加最后一个段
        current_segment.end_address = sorted_data[-1][0]
        current_segment.data_size = len(current_segment.data)
        self.file_info.segments.append(current_segment)
    
    def print_file_info(self):
        """打印文件信息"""
        print("\n=== HEX File Information ===")
        print(f"Filename: {self.file_info.filename}")
        print(f"Total records: {self.file_info.total_records}")
        print(f"Data records: {self.file_info.data_records}")
        print(f"Total data bytes: {self.file_info.total_data_bytes} bytes")
        print(f"Program entry point: 0x{self.file_info.start_address:08X}")
        print(f"End address: 0x{self.file_info.end_address:08X}")
        print(f"Extended linear address: 0x{self.file_info.extended_linear_address:04X}")
        print(f"Extended segment address: 0x{self.file_info.extended_segment_address:04X}")
        print(f"Memory map size: {len(self.file_info.memory_map)} bytes")
        
        # 显示记录类型统计
        type_count = {}
        for record in self.file_info.records:
            type_name = record.record_type.name.replace('_', ' ').title()
            type_count[type_name] = type_count.get(type_name, 0) + 1
        
        print("\nRecord type statistics:")
        for record_type, count in type_count.items():
            print(f"  {record_type}: {count} records")
    
    def print_segment_info(self):
        """打印段信息"""
        print("\n=== Segment Information ===")
        
        for i, segment in enumerate(self.file_info.segments):
            print(f"Segment {i}:")
            print(f"  Start address: 0x{segment.start_address:08X}")
            print(f"  End address: 0x{segment.end_address:08X}")
            print(f"  Size: {segment.data_size} bytes")
            
            # 显示前16字节数据
            print("  First 16 bytes: ", end="")
            for j in range(min(16, len(segment.data))):
                print(f"{segment.data[j]:02X} ", end="")
            print()
    
    def print_sample_data(self):
        """打印示例数据"""
        if not self.file_info.memory_map:
            return
        
        print("\n=== Sample Data Reading ===")
        print(f"First 16 bytes from start address 0x{self.file_info.start_address:08X}:")
        
        # 读取前16字节数据
        data = []
        for i in range(16):
            addr = self.file_info.start_address + i
            data.append(self.file_info.memory_map.get(addr, 0xFF))
        
        # 显示数据
        for i in range(0, len(data), 8):
            addr = self.file_info.start_address + i
            print(f"\n  0x{addr:08X}: ", end="")
            for j in range(min(8, len(data) - i)):
                print(f"{data[i + j]:02X} ", end="")
        print()
    
    def validate_file(self) -> bool:
        """验证文件格式"""
        has_end_record = False
        has_data = False
        
        for record in self.file_info.records:
            if not record.valid:
                print("Found invalid record")
                return False
            
            if record.record_type == HexRecordType.END_OF_FILE:
                has_end_record = True
            
            if record.record_type == HexRecordType.DATA and record.data:
                has_data = True
        
        if not has_data:
            print("Warning: No data records found in file")
        
        if not has_end_record:
            print("Warning: No end of file record found")
        
        return True


def print_usage(program_name: str):
    """打印使用说明"""
    print(f"Usage: {program_name} <hex_file_path_or_single_record>")
    print("Examples:")
    print(f"  {program_name} M6_HMI_LOADER.hex")
    print(f"  {program_name} :04000005080003A547")


def main():
    """主函数"""
    if len(sys.argv) != 2:
        print("Error: Incorrect number of arguments")
        print_usage(sys.argv[0])
        return 1
    
    input_arg = sys.argv[1]
    
    # 检查是否是单行HEX记录
    is_single_line = input_arg.startswith(':')
    
    if not is_single_line:
        # 检查文件扩展名
        if not input_arg.endswith('.hex'):
            print("Error: Input must be either a .hex file or a single HEX record starting with ':'")
            return 1
    
    print(f"Parsing HEX {'record' if is_single_line else 'file'}: {input_arg}")
    
    parser = HexParser()
    
    # 解析
    if is_single_line:
        parse_success = parser.parse_single_record(input_arg)
    else:
        parse_success = parser.parse_file(input_arg)
    
    if not parse_success:
        print(f"Failed to parse {'record' if is_single_line else 'file'}")
        return 1
    
    # 验证文件格式
    if not parser.validate_file():
        print("File format validation failed")
        return 1
    
    # 打印文件信息
    parser.print_file_info()
    
    # 只有完整HEX文件才显示段信息
    if not is_single_line:
        # 分析并打印段信息
        parser.analyze_segments()
        parser.print_segment_info()
    
    # 显示示例数据
    parser.print_sample_data()
    
    print("\nHEX file parsing completed!")
    return 0


if __name__ == "__main__":
    sys.exit(main())
